Avoid tiny numpy arrays in numba: optimize intersections.py (and refactor: "un-scalarize") - #1727

Closes #1726 (sub-issue of #1648)

Overview

Optimizes numba routines in uxarray/grid/intersections.py to avoid constructing many tiny numpy arrays inside numba routines, as discussed in #1648. ASV benchmarks seem to show roughly 2x speedup of intersections algorithms, and no performance degradations. Hard to know for sure how gca_gca intersection performance has been affected because those benchmarks are failing on main right now (see discussion on #1688 for details), but the results here at least seem to be the right order of magnitude (roughly 1 μs; latest successes of these benchmarks on main were roughly 1 to 1.5μs).

Additionally, refactors the intersections.py routines to restore prior behavior (from before #1688 merged) of using 3-vector inputs and outputs where possible. To maintain performance improvements, just use tuples instead of allocating tiny numpy arrays. The discussion in PR 1688 clarified the intent was to restore to the pre-scalarized function interfaces where possible without sacrificing performance.

A few changes are within scope of the original issue but the reasoning is not necessarily obvious directly from inspecting the code diff; clarifying here:

• Previously, gca_gca_intersection output (len(result) == number of intersections) did not match its docstring, which promised shape (2,3) result but filling unused rows with NaNs. This PR updates the implementation to fill unused rows with NaNs, and always return a tuple of 2 length-3 tuples, like ((x1,y1,z1),(x2,y2,z2)), filling with NaNs to represent nonexistent intersection points. (The consistent shape is necessary for the performance optimization here. Numba complains that it "Can't unify return type" if the output length depends on the number of intersections.) The test suite has been updated accordingly, to check for nans instead of len(result).
  • Also updated the docstring of get_number_of_intersections to clarify that it could now be applied to results from gca_gca_intersection too.
• Removes scalarized versions of methods where non-scalarized versions now suffice. I.e., removed _accux_constlat_scalar (use _accux_constlat instead), and _snap_const_lat_endpoint_xy (use _snap_const_lat_endpoint instead). By using tuples instead of tiny numpy arrays, the "non-scalar" versions also avoid allocation costs, so there is no performance-related need to continue maintaining the scalarized versions.
• Similarly, updated arcs.py, removing _on_minor_arc_xyz by moving its logic into on_minor_arc, since the scalarized version wasn't being used anywhere anymore (except in benchmarking suite, which has been updated appropriately to use on_minor_arc). Updated on_minor_arc docstring to clarify inputs don't need to be numpy arrays; tuples of length 3 also work just fine.
• Improved various docstrings, updating docstrings where needed to explain new behavior, but also adding to existing docstrings to clarify meanings and expectations about inputs and outputs, and removing now-obsolete notes about preferring scalarized versions.

EDIT: discussions below revealed the need to incorporate the following changes, too:

• update the geometry_samebody.py and geometry_samebody_gcagca.py benchmarks files to keep the fp64 function performance in line with the optimizations to the corresponding functions from uxarray.
  • While doing that change, I also noticed and fixed: _fp64_try_gca_gca_intersection was using "and" instead of multiplying together 0s and 1s, so fp64 vs accux wasn't actually the only difference between it and intersections.py's _try_gca_gca_intersection.

EDIT: minor expansions of scope (requested in comment below):

• Swapped np.max to np.nanmax in geometry_samebody_gcagca.py:274, to match the const-lat sibling at geometry_samebody.py:296.
• Swapped njit decorators for all _fp64 methods in geometry_samebody and geometry_samebody_gcagca benchmarks to exactly match decorators from intersections.py. This fix may help to resolve Investigate Python thread-scaling of AccuX vs FP64 kernels #1587(?)

PR Checklist

General

• An issue is created and linked
• Added appropriate labels (if your uxarray repo permissions allow it)
• Filled out Overview and Expected Usage (if applicable) sections

Testing & Benchmarking

• There is adequate test coverage of changes from this PR (add new tests if needed)
• If this PR could affect performance, ran ASV benchmarks and confirmed they show expected behavior (add a new benchmark if necessary)

Documentation and Examples

• Docstrings updated with any function changes, and included in all new functions
• User (public) functions added to docs/api.rst; internal (private) function names start with an underscore (_)
• [N/A] If touched any notebook files, cleared the output of all cells before committing
• [N/A] If added new notebook files, put into appropriate directories and referenced in appropriate files

AI Disclosure

AI Usage: GitHub Copilot's inline code suggestions, plus some Claude for numba debugging questions

• I have tested and take responsibility for all AI-generated content in my PR.